{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "entitled-exploration",
   "metadata": {},
   "source": [
    "# Image Clustering\n",
    "\n",
    "This example shows how SentenceTransformer can be used for Image Clustering.\n",
    "\n",
    "As model, we use the [OpenAI CLIP Model](https://github.com/openai/CLIP), which was trained on a large set of images and image alt texts.\n",
    "\n",
    "As a source for fotos, we use the [Unsplash Dataset Lite](https://unsplash.com/data), which contains about 25k images. See the [License](https://unsplash.com/license) about the Unsplash images. \n",
    "\n",
    "We encode all images into vector space and then find high density regions in this vector space, i.e., regions where the images are fairly similar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acoustic-pastor",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mindnlp\n",
    "from sentence_transformers import SentenceTransformer, util\n",
    "from PIL import Image\n",
    "import glob\n",
    "import torch\n",
    "import pickle\n",
    "import zipfile\n",
    "from IPython.display import display\n",
    "from IPython.display import Image as IPImage\n",
    "import os\n",
    "from tqdm.autonotebook import tqdm\n",
    "\n",
    "#First, we load the respective CLIP model\n",
    "model = SentenceTransformer('clip-ViT-B-32')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "major-injury",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Next, we get about 25k images from Unsplash \n",
    "img_folder = 'photos/'\n",
    "if not os.path.exists(img_folder) or len(os.listdir(img_folder)) == 0:\n",
    "    os.makedirs(img_folder, exist_ok=True)\n",
    "    \n",
    "    photo_filename = 'unsplash-25k-photos.zip'\n",
    "    if not os.path.exists(photo_filename):   #Download dataset if does not exist\n",
    "        util.http_get('https://public.ukp.informatik.tu-darmstadt.de/reimers/sentence-transformers/datasets/'+photo_filename, photo_filename)\n",
    "        \n",
    "    #Extract all images\n",
    "    with zipfile.ZipFile(photo_filename, 'r') as zf:\n",
    "        for member in tqdm(zf.infolist(), desc='Extracting'):\n",
    "            zf.extract(member, img_folder)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "earned-mapping",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now, we need to compute the embeddings\n",
    "# To speed things up, we distribute pre-computed embeddings\n",
    "# Otherwise you can also encode the images yourself.\n",
    "# To encode an image, you can use the following code:\n",
    "# from PIL import Image\n",
    "# img_emb = model.encode(Image.open(filepath))\n",
    "\n",
    "use_precomputed_embeddings = True\n",
    "\n",
    "if use_precomputed_embeddings: \n",
    "    emb_filename = 'unsplash-25k-photos-embeddings.pkl'\n",
    "    if not os.path.exists(emb_filename):   #Download dataset if does not exist\n",
    "        util.http_get('http://sbert.net/datasets/'+emb_filename, emb_filename)\n",
    "        \n",
    "    with open(emb_filename, 'rb') as fIn:\n",
    "        img_names, img_emb = pickle.load(fIn)  \n",
    "    print(\"Images:\", len(img_names))\n",
    "else:\n",
    "    img_names = list(glob.glob('unsplash/photos/*.jpg'))\n",
    "    print(\"Images:\", len(img_names))\n",
    "    img_emb = model.encode([Image.open(filepath) for filepath in img_names], batch_size=128, convert_to_tensor=True, show_progress_bar=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "equipped-script",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We have implemented our own, efficient method\n",
    "# to find high density regions in vector space\n",
    "def community_detection(embeddings, threshold, min_community_size=10, init_max_size=1000):\n",
    "    \"\"\"\n",
    "    Function for Fast Community Detection\n",
    "\n",
    "    Finds in the embeddings all communities, i.e. embeddings that are close (closer than threshold).\n",
    "\n",
    "    Returns only communities that are larger than min_community_size. The communities are returned\n",
    "    in decreasing order. The first element in each list is the central point in the community.\n",
    "    \"\"\"\n",
    "\n",
    "    # Compute cosine similarity scores\n",
    "    cos_scores = util.cos_sim(embeddings, embeddings)\n",
    "\n",
    "    # Minimum size for a community\n",
    "    top_k_values, _ = cos_scores.topk(k=min_community_size, largest=True)\n",
    "\n",
    "    # Filter for rows >= min_threshold\n",
    "    extracted_communities = []\n",
    "    for i in range(len(top_k_values)):\n",
    "        if top_k_values[i][-1] >= threshold:\n",
    "            new_cluster = []\n",
    "\n",
    "            # Only check top k most similar entries\n",
    "            top_val_large, top_idx_large = cos_scores[i].topk(k=init_max_size, largest=True)\n",
    "            top_idx_large = top_idx_large.tolist()\n",
    "            top_val_large = top_val_large.tolist()\n",
    "\n",
    "            if top_val_large[-1] < threshold:\n",
    "                for idx, val in zip(top_idx_large, top_val_large):\n",
    "                    if val < threshold:\n",
    "                        break\n",
    "\n",
    "                    new_cluster.append(idx)\n",
    "            else:\n",
    "                # Iterate over all entries (slow)\n",
    "                for idx, val in enumerate(cos_scores[i].tolist()):\n",
    "                    if val >= threshold:\n",
    "                        new_cluster.append(idx)\n",
    "\n",
    "            extracted_communities.append(new_cluster)\n",
    "\n",
    "    # Largest cluster first\n",
    "    extracted_communities = sorted(extracted_communities, key=lambda x: len(x), reverse=True)\n",
    "\n",
    "    # Step 2) Remove overlapping communities\n",
    "    unique_communities = []\n",
    "    extracted_ids = set()\n",
    "\n",
    "    for community in extracted_communities:\n",
    "        add_cluster = True\n",
    "        for idx in community:\n",
    "            if idx in extracted_ids:\n",
    "                add_cluster = False\n",
    "                break\n",
    "\n",
    "        if add_cluster:\n",
    "            unique_communities.append(community)\n",
    "            for idx in community:\n",
    "                extracted_ids.add(idx)\n",
    "\n",
    "    return unique_communities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "oriental-shower",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now we run the clustering algorithm\n",
    "# With the threshold parameter, we define at which threshold we identify\n",
    "# two images as similar. Set the threshold lower, and you will get larger clusters which have \n",
    "# less similar images in it (e.g. black cat images vs. cat images vs. animal images).\n",
    "# With min_community_size, we define that we only want to have clusters of a certain minimal size\n",
    "clusters = community_detection(img_emb, threshold=0.9, min_community_size=10)\n",
    "print(\"Total number of clusters:\", len(clusters))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "buried-botswana",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now we output the first 10 (largest) clusters\n",
    "for cluster in clusters[0:10]:\n",
    "    print(\"\\n\\nCluster size:\", len(cluster))\n",
    "    \n",
    "    #Output 3 images\n",
    "    for idx in cluster[0:3]:\n",
    "        display(IPImage(os.path.join(img_folder, img_names[idx]), width=200))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mindspore",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
